import cv2
import numpy as np
import os
import argparse
import cv2
from os import path as osp

# 设置中文显示
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Microsoft YaHei", "Arial Unicode MS"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


def load_image(image_path):
    """加载图像并转换为RGB格式"""
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法找到图像: {image_path}")
    # 转换为RGB格式以便matplotlib正确显示
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


def calculate_difference(img1, img2, normalize=True):
    """计算两张图像之间的差异"""
    # 确保图像尺寸相同
    if img1.shape != img2.shape:
        h, w = min(img1.shape[0], img2.shape[0]), min(img1.shape[1], img2.shape[1])
        img1 = img1[:h, :w]
        img2 = img2[:h, :w]

    # 计算像素级差异
    diff = cv2.absdiff(img1, img2)

    # 如果需要归一化差异以便更好地可视化
    if normalize:
        diff = cv2.normalize(diff, None, 0, 255, cv2.NORM_MINMAX)

    return diff


def enhance_difference(diff, factor=2.0):
    """增强差异图像的对比度，使差异更明显"""
    # 将差异乘以一个因子增强对比度
    enhanced_diff = np.clip(diff * factor, 0, 255).astype(np.uint8)
    return enhanced_diff


def visualize_diff_gt_sr(sr_path, gt_path, output_path=None, enhance_factor=2.0):
    """可视化超分辨率图像和原始高质量图像(GT)之间的差异"""
    # 加载图像
    sr_img = load_image(sr_path)
    gt_img = load_image(gt_path)

    # 创建画布
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    # 显示超分辨率图像
    axes[0].imshow(sr_img)
    axes[0].set_title('超分辨率图像')
    axes[0].axis('off')

    # 显示原始高质量图像(GT)
    axes[1].imshow(gt_img)
    axes[1].set_title('原始高质量图像(GT)')
    axes[1].axis('off')

    # 计算GT和SR的差异并增强显示
    diff_gt_sr = calculate_difference(gt_img, sr_img)
    enhanced_diff_gt_sr = enhance_difference(diff_gt_sr, enhance_factor)

    # 显示差异图
    im = axes[2].imshow(enhanced_diff_gt_sr, cmap='jet')
    axes[2].set_title(f'GT与SR差异图(增强系数: {enhance_factor})')
    axes[2].axis('off')

    # 添加颜色条
    cbar = fig.colorbar(im, ax=axes[2], fraction=0.046, pad=0.04)
    cbar.set_label('差异强度')

    fig.tight_layout()

    # 如果提供了输出路径，保存图像
    if output_path:
        # 确保输出目录存在
        os.makedirs(osp.dirname(osp.abspath(output_path)), exist_ok=True)
        fig.savefig(output_path, dpi=300, bbox_inches='tight')
        print(f"GT与SR差异可视化已保存至: {output_path}")

    # 显示图像
    plt.show()


def batch_visualize_diff_gt_sr(results_dir, gt_dir, output_dir=None, enhance_factor=2.0, num_samples=5):
    """批量可视化超分辨率图像和原始高质量图像(GT)之间的差异"""
    # 获取所有超分辨率图像
    sr_files = [f for f in os.listdir(results_dir) if f.endswith('.png')]

    # 限制样本数量
    sr_files = sr_files[:num_samples]

    if not sr_files:
        print("未找到超分辨率图像。")
        return

    # 确保输出目录存在
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    # 逐个处理图像
    for sr_file in sr_files:
        # 获取图像名称（不包含扩展名）
        img_name = osp.splitext(sr_file)[0]
        print(f"处理图像: {img_name}")

        # 构建文件路径
        sr_path = osp.join(results_dir, sr_file)

        # 尝试找到对应的GT图像
        # 假设原始文件名在超分辨率文件名中
        # 这可能需要根据实际的命名规则调整
        gt_file = None
        for f in os.listdir(gt_dir):
            # 尝试不同的匹配策略以提高匹配成功率
            if (img_name.split('_')[0] in f or f.split('_')[0] in img_name or \
                img_name.replace('SR', '') in f or f in img_name.replace('SR', '')):
                gt_file = f
                break

        if not gt_file:
            print(f"警告: 未找到与 {sr_file} 对应的GT图像。")
            continue

        gt_path = osp.join(gt_dir, gt_file)

        # 设置输出路径
        output_path = None
        if output_dir:
            output_path = osp.join(output_dir, f"{img_name}_gt_sr_diff.png")

        # 可视化差异
        visualize_diff_gt_sr(sr_path, gt_path, output_path, enhance_factor)


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="可视化超分辨率图像和原始高质量图像(GT)之间的差异")

    # 设置两种模式：单图像模式和批量模式
    parser.add_argument('--mode', type=str, default='single', choices=['single', 'batch'],
                        help='可视化模式：single（单图像）或 batch（批量）')

    # 单图像模式参数
    parser.add_argument('--sr_path', type=str, help='超分辨率图像路径')
    parser.add_argument('--gt_path', type=str, help='原始高质量图像路径')
    parser.add_argument('--output_path', type=str, default=None, help='输出图像保存路径（可选）')

    # 批量模式参数
    parser.add_argument('--results_dir', type=str, help='超分辨率结果目录路径')
    parser.add_argument('--gt_dir', type=str, help='原始高质量图像目录路径')
    parser.add_argument('--output_dir', type=str, default=None, help='输出图像保存目录（可选）')
    parser.add_argument('--num_samples', type=int, default=5, help='批量模式下处理的样本数量')

    # 通用参数
    parser.add_argument('--enhance_factor', type=float, default=2.0, help='差异图增强系数')

    args = parser.parse_args()

    # 根据模式执行不同的操作
    if args.mode == 'single':
        # 检查必要参数
        if not args.sr_path or not args.gt_path:
            parser.error("单图像模式需要提供 --sr_path 和 --gt_path 参数")

        # 可视化单张图像的差异
        visualize_diff_gt_sr(args.sr_path, args.gt_path, args.output_path, args.enhance_factor)
    else:
        # 检查必要参数
        if not args.results_dir or not args.gt_dir:
            parser.error("批量模式需要提供 --results_dir 和 --gt_dir 参数")

        # 批量可视化图像差异
        batch_visualize_diff_gt_sr(args.results_dir, args.gt_dir, args.output_dir,
                                  args.enhance_factor, args.num_samples)